{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "q4WF3l23pumU"
      },
      "source": [
        "##### Copyright 2018 The AdaNet Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "Kic2quJWppmx"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "aL7SpaKdirqG"
      },
      "source": [
        "# AdaNet on TPU"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "i5s1gsS1bOuB"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/adanet_tpu.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/adanet/blob/master/adanet/examples/tutorials/adanet_tpu.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "ySKpPtPrmaCx"
      },
      "source": [
        "AdaNet supports training on Google's custom machine learning accelerators known\n",
        "as Tensor Processing Units (TPU). Conveniently, we provide `adanet.TPUEstimator`\n",
        "which handles TPU support behind the scenes. There are only a few minor changes\n",
        "needed to switch from `adanet.Estimator` to `adanet.TPUEstimator`. We highlight\n",
        "the necessary changes in this tutorial.\n",
        "\n",
        "If the reader is not familiar with AdaNet, it is reccommended they take a look\n",
        "at\n",
        "[The AdaNet Objective](https://colab.sandbox.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/adanet_objective.ipynb)\n",
        "and in particular\n",
        "[Customizing AdaNet](https://colab.sandbox.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/adanet_objective.ipynb)\n",
        "as this tutorial builds upon the latter.\n",
        "\n",
        "**NOTE: you must provide a valid GCS bucket to use TPUEstimator.**\n",
        "\n",
        "To begin, we import the necessary packages, obtain the Colab's TPU master\n",
        "address, and give the TPU permissions to write to our GCS Bucket. Follow the\n",
        "instructions\n",
        "[here](https://colab.sandbox.google.com/notebooks/tpu.ipynb#scrollTo=_pQCOmISAQBu)\n",
        "to connect to a Colab TPU runtime.\n",
        "\n"]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "VfGCvk-NLY35"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "# If you're running this in Colab, first install the adanet package:\n",
        "!pip install adanet"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "x_3b6xx2s6B9"
      },
      "outputs": [],
      "source": [
        "from __future__ import absolute_import\n",
        "from __future__ import division\n",
        "from __future__ import print_function\n",
        "\n",
        "import functools\n",
        "import json\n",
        "import os\n",
        "import six\n",
        "import time\n",
        "\n",
        "import adanet\n",
        "from google.colab import auth\n",
        "import tensorflow as tf\n",
        "\n",
        "BUCKET = ''  #@param {type: 'string'}\n",
        "MODEL_DIR = 'gs://{}/{}'.format(\n",
        "    BUCKET, time.strftime('adanet-tpu-estimator/%Y-%m-%d-%H-%M-%S'))\n",
        "\n",
        "MASTER = ''\n",
        "if 'COLAB_TPU_ADDR' in os.environ:\n",
        "  auth.authenticate_user()\n",
        "\n",
        "  MASTER = 'grpc://' + os.environ['COLAB_TPU_ADDR']\n",
        "\n",
        "  # Authenticate TPU to use GCS Bucket.\n",
        "  with tf.Session(MASTER) as sess:\n",
        "    with open('/content/adc.json', 'r') as file_:\n",
        "      auth_info = json.load(file_)\n",
        "    tf.contrib.cloud.configure_gcs(sess, credentials=auth_info)\n",
        "\n",
        "\n",
        "# The random seed to use.\n",
        "RANDOM_SEED = 42"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "5_hRtdchqRZb"
      },
      "source": [
        "## Fashion MNIST\n",
        "\n",
        "We focus again on the Fashion MNIST dataset and download the data via Keras."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "uYklOnPJ4h7g",
        "outputId": "85206a54-7213-4623-c915-e38e6f0c7e0f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-labels-idx1-ubyte.gz\n",
            "32768/29515 [=================================] - 0s 0us/step\n",
            "40960/29515 [=========================================] - 0s 0us/step\n",
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/train-images-idx3-ubyte.gz\n",
            "26427392/26421880 [==============================] - 0s 0us/step\n",
            "26435584/26421880 [==============================] - 0s 0us/step\n",
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-labels-idx1-ubyte.gz\n",
            "16384/5148 [===============================================================================================] - 0s 0us/step\n",
            "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/t10k-images-idx3-ubyte.gz\n",
            "4423680/4422102 [==============================] - 0s 0us/step\n",
            "4431872/4422102 [==============================] - 0s 0us/step\n"
          ]
        }
      ],
      "source": [
        "(x_train, y_train), (x_test, y_test) = (\n",
        "    tf.keras.datasets.fashion_mnist.load_data())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "tECo5dFd4QCa"
      },
      "source": [
        "## `input_fn` Changes\n",
        "\n",
        "There are two minor changes we must make to `input_fn` to support running on\n",
        "TPU:\n",
        "\n",
        "1.  TPUs dynamically shard the input data depending on the number of cores used.\n",
        "    Because of this, we augment `input_fn` to take a dictionary `params`\n",
        "    argument. When running on TPU, `params` contains a `batch_size` field with\n",
        "    the appropriate batch size.\n",
        "\n",
        "1.  Once the input is batched, we drop the last batch if it is smaller than\n",
        "    `batch_size`. This can simply be done by specifying `drop_remainder=True` to\n",
        "    the\n",
        "    [`tf.data.DataSet.batch()`](https://www.tensorflow.org/api_docs/python/tf/data/Dataset#batch)\n",
        "    function. It is important to specify this option since TPUs do not support\n",
        "    dynamic shapes. Note that we only drop the remainder batch during training\n",
        "    since evaluation is still done on the CPU."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "gxTAoIXwsTH7"
      },
      "outputs": [],
      "source": [
        "FEATURES_KEY = \"images\"\n",
        "\n",
        "\n",
        "def generator(images, labels):\n",
        "  \"\"\"Returns a generator that returns image-label pairs.\"\"\"\n",
        "\n",
        "  def _gen():\n",
        "    for image, label in zip(images, labels):\n",
        "      yield image, label\n",
        "\n",
        "  return _gen\n",
        "\n",
        "\n",
        "def preprocess_image(image, label):\n",
        "  \"\"\"Preprocesses an image for an `Estimator`.\"\"\"\n",
        "  image = image / 255.\n",
        "  image = tf.reshape(image, [28, 28, 1])\n",
        "  features = {FEATURES_KEY: image}\n",
        "  return features, label\n",
        "\n",
        "\n",
        "def input_fn(partition, training, batch_size):\n",
        "  \"\"\"Generate an input_fn for the Estimator.\"\"\"\n",
        "\n",
        "  def _input_fn(params):  # TPU: specify `params` argument.\n",
        "\n",
        "    # TPU: get the TPU set `batch_size`, if available.\n",
        "    batch_size_ = params.get(\"batch_size\", batch_size)\n",
        "\n",
        "    if partition == \"train\":\n",
        "      dataset = tf.data.Dataset.from_generator(\n",
        "          generator(x_train, y_train), (tf.float32, tf.int32), ((28, 28), ()))\n",
        "    elif partition == \"predict\":\n",
        "      dataset = tf.data.Dataset.from_generator(\n",
        "          generator(x_test[:10], y_test[:10]), (tf.float32, tf.int32),\n",
        "          ((28, 28), ()))\n",
        "    else:\n",
        "      dataset = tf.data.Dataset.from_generator(\n",
        "          generator(x_test, y_test), (tf.float32, tf.int32), ((28, 28), ()))\n",
        "\n",
        "    if training:\n",
        "      dataset = dataset.shuffle(10 * batch_size_, seed=RANDOM_SEED).repeat()\n",
        "\n",
        "    # TPU: drop the remainder batch when training on TPU.\n",
        "    dataset = dataset.map(preprocess_image).batch(\n",
        "        batch_size_, drop_remainder=training)\n",
        "    iterator = dataset.make_one_shot_iterator()\n",
        "    features, labels = iterator.get_next()\n",
        "    return features, labels\n",
        "\n",
        "  return _input_fn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "D3IE6-9vFVlg"
      },
      "source": [
        "## `model_fn` Changes\n",
        "\n",
        "We use a similar CNN architecture as used in the\n",
        "[Customizing AdaNet](https://colab.sandbox.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet.ipynb)\n",
        "tutorial. The only TPU specific change we need to make is wrap the `optimizer`\n",
        "in a\n",
        "[`tf.contrib.tpu.CrossShardOptimizer`](https://www.google.com/search?q=cross+shard+optimizer\u0026oq=cross+shard+optimizer\u0026aqs=chrome.0.0j69i57.2391j0j7\u0026sourceid=chrome\u0026ie=UTF-8)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "id": "IsYJ97tRwBkt"
      },
      "outputs": [],
      "source": [
        "#@title Define the Builder and Generator\n",
        "class SimpleCNNBuilder(adanet.subnetwork.Builder):\n",
        "  \"\"\"Builds a CNN subnetwork for AdaNet.\"\"\"\n",
        "\n",
        "  def __init__(self, learning_rate, max_iteration_steps, seed):\n",
        "    \"\"\"Initializes a `SimpleCNNBuilder`.\n",
        "\n",
        "    Args:\n",
        "      learning_rate: The float learning rate to use.\n",
        "      max_iteration_steps: The number of steps per iteration.\n",
        "      seed: The random seed.\n",
        "\n",
        "    Returns:\n",
        "      An instance of `SimpleCNNBuilder`.\n",
        "    \"\"\"\n",
        "    self._learning_rate = learning_rate\n",
        "    self._max_iteration_steps = max_iteration_steps\n",
        "    self._seed = seed\n",
        "\n",
        "  def build_subnetwork(self,\n",
        "                       features,\n",
        "                       logits_dimension,\n",
        "                       training,\n",
        "                       iteration_step,\n",
        "                       summary,\n",
        "                       previous_ensemble=None):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    images = list(features.values())[0]\n",
        "    kernel_initializer = tf.keras.initializers.he_normal(seed=self._seed)\n",
        "    x = tf.keras.layers.Conv2D(\n",
        "        filters=16,\n",
        "        kernel_size=3,\n",
        "        padding=\"same\",\n",
        "        activation=\"relu\",\n",
        "        kernel_initializer=kernel_initializer)(\n",
        "            images)\n",
        "    x = tf.keras.layers.MaxPool2D(pool_size=2, strides=2)(x)\n",
        "    x = tf.keras.layers.Flatten()(x)\n",
        "    x = tf.keras.layers.Dense(\n",
        "        units=64, activation=\"relu\", kernel_initializer=kernel_initializer)(\n",
        "            x)\n",
        "\n",
        "    logits = tf.keras.layers.Dense(\n",
        "        units=10, activation=None, kernel_initializer=kernel_initializer)(\n",
        "            x)\n",
        "\n",
        "    complexity = tf.constant(1)\n",
        "\n",
        "    return adanet.Subnetwork(\n",
        "        last_layer=x,\n",
        "        logits=logits,\n",
        "        complexity=complexity,\n",
        "        persisted_tensors={})\n",
        "\n",
        "  def build_subnetwork_train_op(self,\n",
        "                                subnetwork,\n",
        "                                loss,\n",
        "                                var_list,\n",
        "                                labels,\n",
        "                                iteration_step,\n",
        "                                summary,\n",
        "                                previous_ensemble=None):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "\n",
        "    learning_rate = tf.train.cosine_decay(\n",
        "        learning_rate=self._learning_rate,\n",
        "        global_step=iteration_step,\n",
        "        decay_steps=self._max_iteration_steps)\n",
        "    optimizer = tf.train.MomentumOptimizer(learning_rate, .9)\n",
        "    # TPU: wrap the optimizer in a CrossShardOptimizer.\n",
        "    optimizer = tf.contrib.tpu.CrossShardOptimizer(optimizer)\n",
        "    return optimizer.minimize(loss=loss, var_list=var_list)\n",
        "\n",
        "  def build_mixture_weights_train_op(self, loss, var_list, logits, labels,\n",
        "                                     iteration_step, summary):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    return tf.no_op(\"mixture_weights_train_op\")\n",
        "\n",
        "  @property\n",
        "  def name(self):\n",
        "    \"\"\"See `adanet.subnetwork.Builder`.\"\"\"\n",
        "    return \"simple_cnn\"\n",
        "\n",
        "\n",
        "class SimpleCNNGenerator(adanet.subnetwork.Generator):\n",
        "  \"\"\"Generates a `SimpleCNN` at each iteration.\"\"\"\n",
        "\n",
        "  def __init__(self, learning_rate, max_iteration_steps, seed=None):\n",
        "    \"\"\"Initializes a `Generator` that builds `SimpleCNNs`.\n",
        "\n",
        "    Args:\n",
        "      learning_rate: The float learning rate to use.\n",
        "      max_iteration_steps: The number of steps per iteration.\n",
        "      seed: The random seed.\n",
        "\n",
        "    Returns:\n",
        "      An instance of `Generator`.\n",
        "    \"\"\"\n",
        "    self._seed = seed\n",
        "    self._dnn_builder_fn = functools.partial(\n",
        "        SimpleCNNBuilder,\n",
        "        learning_rate=learning_rate,\n",
        "        max_iteration_steps=max_iteration_steps)\n",
        "\n",
        "  def generate_candidates(self, previous_ensemble, iteration_number,\n",
        "                          previous_ensemble_reports, all_reports):\n",
        "    \"\"\"See `adanet.subnetwork.Generator`.\"\"\"\n",
        "    seed = self._seed\n",
        "    # Change the seed according to the iteration so that each subnetwork\n",
        "    # learns something different.\n",
        "    if seed is not None:\n",
        "      seed += iteration_number\n",
        "    return [self._dnn_builder_fn(seed=seed)]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fhu2zpf9faIB"
      },
      "source": [
        "## Launch TensorBoard\n",
        "\n",
        "Let's run [TensorBoard](https://www.tensorflow.org/guide/summaries_and_tensorboard) to visualize model training over time. We'll use [ngrok](https://ngrok.com/) to tunnel traffic to localhost.\n",
        "\n",
        "*The instructions for setting up Tensorboard were obtained from https://www.dlology.com/blog/quick-guide-to-run-tensorboard-in-google-colab/*\n",
        "\n",
        "Run the next cells and follow the link to see the TensorBoard in a new tab."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "S1UcP5yeaDz9"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "\n",
        "get_ipython().system_raw(\n",
        "    'tensorboard --logdir {} --host 0.0.0.0 --port 6006 \u0026'\n",
        "    .format(MODEL_DIR)\n",
        ")\n",
        "\n",
        "# Install ngrok binary.\n",
        "! wget https://bin.equinox.io/c/4VmDzA7iaHb/ngrok-stable-linux-amd64.zip\n",
        "! unzip ngrok-stable-linux-amd64.zip\n",
        "\n",
        "print(\"Follow this link to open TensorBoard in a new tab.\")\n",
        "get_ipython().system_raw('./ngrok http 6006 \u0026')\n",
        "! curl -s http://localhost:4040/api/tunnels | python3 -c \\\n",
        "    \"import sys, json; print(json.load(sys.stdin)['tunnels'][0]['public_url'])\"\n",
        "\n"]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "8sdvharsLJ1T"
      },
      "source": [
        "## Using `adanet.TPUEstimator` to Train and Evaluate\n",
        "\n",
        "Finally, we switch from `adanet.Estimator` to `adanet.TPUEstimator`. There are\n",
        "two last changes needed:\n",
        "\n",
        "1.  Update the `RunConfig` to be a\n",
        "    [`tf.contrib.tpu.RunConfig`](https://www.tensorflow.org/api_docs/python/tf/contrib/tpu/RunConfig).\n",
        "    We supply the TPU `master` address and set `iterations_per_loop=200`. This\n",
        "    choice is fairly arbitrary in our case. A good practice is to set it to the\n",
        "    number of steps in between summary writes and metric evals.\n",
        "1.  Finally, we specify the `use_tpu` and `batch_size` parameters\n",
        "    `adanet.TPUEstimator`.\n",
        "\n",
        "There is an important thing to note about the `batch_size`: each TPU chip\n",
        "consists of 2 cores with 4 shards each. In the\n",
        "[Customizing AdaNet](https://colab.sandbox.google.com/github/tensorflow/adanet/blob/master/adanet/examples/tutorials/customizing_adanet.ipynb)\n",
        "tutorial, a `batch_size` of 64 was used. To be consistent we use the same\n",
        "`batch_size` per shard and drop the number of training steps accordingly. In\n",
        "other words, since we're running on one TPU we set `batch_size=64*8=512` and\n",
        "`train_steps=1000`. In the ideal case, since we drop the `train_steps` by 5x,\n",
        "this means we're **training 5x faster!**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 0,
      "metadata": {
        "cellView": "both",
        "colab": {
          "height": 53,
          "test": {
            "output": "ignore",
            "timeout": 900
          }
        },
        "colab_type": "code",
        "executionInfo": {
          "elapsed": 78671,
          "status": "ok",
          "timestamp": 1545239915199,
          "user": {
            "displayName": "",
            "photoUrl": "",
            "userId": ""
          },
          "user_tz": 300
        },
        "id": "-Fhi1SjkzVBt",
        "outputId": "4e56038d-5424-4170-ee40-460811ee93fd"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Accuracy: 0.8913\n",
            "Loss: 0.298405\n"
          ]
        }
      ],
      "source": [
        "#@title AdaNet Parameters\n",
        "LEARNING_RATE = 0.25  #@param {type:\"number\"}\n",
        "TRAIN_STEPS = 1000  #@param {type:\"integer\"}\n",
        "BATCH_SIZE = 512  #@param {type:\"integer\"}\n",
        "ADANET_ITERATIONS = 2  #@param {type:\"integer\"}\n",
        "\n",
        "# TPU: switch `tf.estimator.RunConfig` to `tf.contrib.tpu.RunConfig`.\n",
        "# The main required changes are specifying `tpu_config` and `master`.\n",
        "config = tf.contrib.tpu.RunConfig(\n",
        "    tpu_config=tf.contrib.tpu.TPUConfig(iterations_per_loop=200),\n",
        "    master=MASTER,\n",
        "    save_checkpoints_steps=200,\n",
        "    save_summary_steps=200,\n",
        "    tf_random_seed=RANDOM_SEED)\n",
        "\n",
        "head = tf.contrib.estimator.multi_class_head(\n",
        "    n_classes=10, loss_reduction=tf.losses.Reduction.SUM_OVER_BATCH_SIZE)\n",
        "max_iteration_steps = TRAIN_STEPS // ADANET_ITERATIONS\n",
        "# TPU: switch `adanet.Estimator` to `adanet.TPUEstimator`.\n",
        "try:\n",
        "  estimator = adanet.TPUEstimator(\n",
        "      head=head,\n",
        "      subnetwork_generator=SimpleCNNGenerator(\n",
        "          learning_rate=LEARNING_RATE,\n",
        "          max_iteration_steps=max_iteration_steps,\n",
        "          seed=RANDOM_SEED),\n",
        "      max_iteration_steps=max_iteration_steps,\n",
        "      evaluator=adanet.Evaluator(\n",
        "          input_fn=input_fn(\"train\", training=False, batch_size=BATCH_SIZE),\n",
        "          steps=None),\n",
        "      adanet_loss_decay=.99,\n",
        "      config=config,\n",
        "      model_dir=MODEL_DIR,\n",
        "      # TPU: specify `use_tpu` and the batch_size parameters.\n",
        "      use_tpu=True,\n",
        "      train_batch_size=BATCH_SIZE,\n",
        "      eval_batch_size=32)\n",
        "except tf.errors.InvalidArgumentError as e:\n",
        "  six.raise_from(\n",
        "      Exception(\n",
        "          \"Invalid GCS Bucket: you must provide a valid GCS bucket in the \"\n",
        "          \"`BUCKET` form field of the first cell.\"), e)\n",
        "\n",
        "results, _ = tf.estimator.train_and_evaluate(\n",
        "    estimator,\n",
        "    train_spec=tf.estimator.TrainSpec(\n",
        "        input_fn=input_fn(\"train\", training=True, batch_size=BATCH_SIZE),\n",
        "        max_steps=TRAIN_STEPS),\n",
        "    eval_spec=tf.estimator.EvalSpec(\n",
        "        input_fn=input_fn(\"test\", training=False, batch_size=BATCH_SIZE),\n",
        "        steps=None,\n",
        "        start_delay_secs=1,\n",
        "        throttle_secs=1,\n",
        "    ))\n",
        "\n",
        "print(\"Accuracy:\", results[\"accuracy\"])\n",
        "print(\"Loss:\", results[\"average_loss\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TKhCzP65hGyS"
      },
      "source": [
        "## Conclusion\n",
        "\n",
        "That was easy! With very few changes we were able to transform our original\n",
        "estimator into one which can harness the power of TPUs."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
            "name": "adanet_tpu.ipynb",
      "provenance": [],
      "version": "0.3.2"
    },
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
